package com.small.sentinel.slots.flow.controller;

import com.small.sentinel.node.Node;
import com.small.sentinel.slots.block.RuleConstant;

public class DefaultController implements TrafficShapingController{

    private static final int DEFAULT_AVG_USED_TOKENS = 0;

    /**
     * 流量控制阈值计数
     */
    private double count;

    /**
     * 0:线程数 1:qps
     */
    private int grade;

    public DefaultController(double count, int grade) {
        this.count = count;
        this.grade = grade;
    }

    @Override
    public boolean canPass(Node node, int acquireCount) {
        int curCount = avgUsedTokens(node);
        if (curCount + acquireCount > count) {
            return false;
        }
        return true;
    }


    /**
     * 平均使用令牌数
     * @param node
     * @return
     */
    private int avgUsedTokens(Node node) {
        if (node == null) {
            return DEFAULT_AVG_USED_TOKENS;
        }
        return grade == RuleConstant.FLOW_GRADE_THREAD ? node.curThreadNum() : (int)(node.passQps());
    }
}
